# Licensed under the MIT license.

"""
rStar推理系统的MCTS实现

本模块实现了基于MCTS的数学推理系统，包括：
1. Generator类 - 负责生成各种类型的推理节点
2. Reasoning_MCTS_Node类 - 专门用于推理任务的MCTS节点
3. search_for_answers函数 - 执行完整的MCTS搜索过程

主要功能：
- 支持多种推理动作：直接回答、子问题分解、一步思考、问题重述等
- 实现潜在分数评估，提高解决方案选择的准确性
- 提供完整的推理轨迹记录和可视化
- 支持多种推理策略的组合使用

核心算法：
- 使用MCTS算法探索推理空间
- 通过UCT公式平衡探索与利用
- 基于置信度和潜在分数选择最佳解决方案

应用场景：
- 数学问题求解
- 逻辑推理
- 复杂问题的分步解决
"""

import sys

sys.path.append(".")

# 基础库导入
import numpy as np  # 数值计算和数组操作
import os  # 操作系统接口，文件路径操作
import random  # 随机数生成
import json  # JSON数据处理
import math  # 数学函数
import wandb  # 实验跟踪和可视化
from tqdm import trange  # 进度条显示
from typing import List, Dict, Tuple  # 类型提示
from copy import deepcopy  # 深拷贝对象

# 可选库导入
try:
    from rapidfuzz import fuzz, process  # 快速字符串匹配和相似度计算
except:
    pass

# 项目内部模块导入
from models.IO_System import IO_System  # 模型输入输出系统
from common.utils import read_txt, read_json  # 文件读取工具函数
from eval_src.Evaluator import Evaluator, GSM8KEvaluator  # 答案评估器
from run_src.MCTS_backbone import MCTS_Searcher, MCTS_Node  # MCTS搜索算法核心
from run_src.rstar_utils import (  # rStar推理系统工具函数
    Node_Type,  # 节点类型枚举
    GeneratorError,  # 生成器异常类
    reach_terminal_subquestion,  # 判断是否到达终端子问题
    reach_terminal_ost_step,  # 判断是否到达终端思考步骤
    concat_subqs_and_subas,  # 拼接子问题和子答案
    concat_ost_steps,  # 拼接一步思考步骤
    concat_subqs_subas_as_ost_steps,  # 将子问题和子答案转换为思考步骤
    make_hint,  # 生成提示信息
    make_response_prefix,  # 生成响应前缀
    split_user_question,  # 分割用户问题
    print_tree_from_root,  # 从根节点打印搜索树
    find_valid_solution_nodes,  # 查找有效解决方案节点
    find_best_solution,  # 查找最佳解决方案
    stochastic_find_best_solution,  # 随机查找最佳解决方案
)


def verbose_print(s: str, verbose: bool):
    """
    条件打印函数，仅在verbose模式下输出信息
    
    Args:
        s (str): 待打印的字符串
        verbose (bool): 是否处于详细输出模式
    """
    if verbose:
        print(s)


class Generator:
    """
    推理节点生成器
    
    负责生成MCTS推理树中的各种类型的子节点，包括：
    - 直接答案节点：基于当前状态直接回答问题
    - 子问题节点：将复杂问题分解为更简单的子问题
    - 重新回答节点：对已有子问题提供新的答案
    - 一步思考节点：提供单步推理思路
    - 问题重述节点：重新表述原问题以获得不同视角
    
    核心功能：
    - 使用大语言模型生成各种推理内容
    - 支持多种推理策略的组合使用
    - 实现潜在分数评估以提高解决方案质量
    - 提供缓存机制以提高效率
    """

    def __init__(self, args, tokenizer, model, evaluator: Evaluator) -> None:
        """
        初始化推理节点生成器
        
        Args:
            args: 命令行参数对象，包含各种配置选项
            tokenizer: 分词器
            model: 语言模型
            evaluator: 评估器实例，用于答案提取和比较
        """
        # 初始化IO系统和评估器
        self.io = IO_System(args, tokenizer, model)
        self.evaluator = evaluator

        # 生成参数配置
        self.num_subquestions = args.num_subquestions      # 生成子问题的数量
        self.num_a1_steps = args.num_a1_steps              # 生成一步思考的数量
        
        # ====== 投票参数配置：两种不同的投票策略 ======
        self.num_votes = args.num_votes                    # 【中间节点投票数量】用于非终端子问题的答案选择
        # 使用场景：
        # 1. A3行为：为非终端子问题生成答案候选时使用
        # 2. A4行为：重新回答子问题时使用  
        # 3. 潜在分数评估：评估中间节点的潜在价值时使用
        # 目的：在推理过程中快速选择相对可靠的中间答案，平衡质量和效率
        
        self.mcts_num_last_votes = args.mcts_num_last_votes  # 【终端节点投票数量】用于终端子问题和直接答案的选择
        # 使用场景：
        # 1. 终端子问题：当reach_terminal_subquestion()返回True时使用
        # 2. 直接答案生成：generate_direct_answers()中使用
        # 目的：在推理的最后阶段使用更多候选来确保最终答案的高质量和可靠性
        
        self.max_tokens = args.max_tokens                  # 最大生成token数
        self.enable_potential_score = args.enable_potential_score  # 是否启用潜在分数评估

        # 加载问题索引配置
        with open(args.decompose_template_path, "r") as f:
            decompose_template = json.load(f)
            self.question_index = decompose_template["index"]

        # 加载基础提示模板
        self.decompose_prompt = read_txt(args.decompose_prompt_path)        # 问题分解提示
        self.fewshot_cot_prompt = read_txt(args.fewshot_cot_prompt_path)    # 少样本思维链提示
        self.fewshot_cot_config = read_json(args.fewshot_cot_config_path)   # 思维链配置

        # 加载一步思考相关提示（A1动作）
        if not args.disable_a1:  # A1: Propose an one-step thought.
            self.fewshot_ost_prompt = read_txt(args.fewshot_ost_prompt_path)
            self.fewshot_ost_config = read_json(args.fewshot_ost_config_path)

        # 加载问题重述相关提示（A5动作）
        if not args.disable_a5:  # A5: Rephrase the question/sub-question.
            self.rephrasing_prompt_template = read_txt(args.rephrasing_prompt_template_path)
            self.decompose_prompt_rephrased = read_txt(args.decompose_prompt_rephrased_path)
            self.fewshot_cot_prompt_rephrased = read_txt(args.fewshot_cot_prompt_rephrased_path)
            self.fewshot_ost_prompt_rephrased = read_txt(args.fewshot_ost_prompt_rephrased_path)

    def _extract_from_cache(self, subquestion_list: List[str]):
        """
        从推理缓存中提取相似子问题的答案
        
        使用字符串相似度匹配技术从缓存中查找与当前子问题列表相似的已解决问题，
        并根据相似度和历史解决质量对问题进行分类处理。
        
        【核心功能】：
        1. 字符串相似度匹配：使用rapidfuzz库计算问题间的相似度
        2. 质量分层处理：根据历史解决质量将匹配结果分为高分和低分两类
        3. 缓存命中优化：对于高质量且完全匹配的问题直接复用答案
        
        【分类策略】：
        - 高分问题(score >= 0.9)：历史解决质量高，可直接复用答案
        - 低分问题(score < 0.9)：历史解决质量一般，提供候选答案列表供选择
        - 未匹配问题：缓存中无相似问题，需要重新生成答案
        
        Args:
            subquestion_list (List[str]): 待匹配的子问题列表
            
        Returns:
            dict: 包含以下键的字典：
                - high_score_questions: 高质量匹配的问题列表
                - selected_answers: 高质量问题对应的最佳答案列表
                - values: 高质量问题的历史解决质量分数
                - low_score_questions: 低质量匹配的问题列表
                - low_score_values: 低质量问题的历史解决质量分数
                - low_score_answers_list: 低质量问题的候选答案列表
                - unmatched_questions: 缓存中无匹配的问题列表
        """
        high_score_questions = []
        selected_answers = []
        values = []
        low_score_questions = []
        low_score_values = []
        low_score_answers_list = []
        unmatched_questions = []

        for subquestion in subquestion_list:
            best_match = process.extractOne(subquestion, self.reasoning_cache.keys(), scorer=fuzz.ratio)

            if best_match:
                best_question, best_score = best_match[0], best_match[1]
                similarity = best_score / 100
                cache_entry = self.reasoning_cache[best_question]
                score = cache_entry["score"]
                if similarity == 1:
                    if score >= 0.9:
                        high_score_questions.append(best_question)
                        selected_answers.append(cache_entry["selected_answer"])
                        values.append(score)
                    else:
                        low_score_questions.append(best_question)
                        low_score_values.append(score)
                        low_score_answers_list.append(cache_entry["answer_list"])
                else:
                    unmatched_questions.append(subquestion)
            else:
                unmatched_questions.append(subquestion)

        return {
            "high_score_questions": high_score_questions,
            "selected_answers": selected_answers,  # most likely answer corresponding to each subquestion
            "values": values,
            "low_score_questions": low_score_questions,
            "low_score_values": low_score_values,
            "low_score_answers_list": low_score_answers_list,
            "unmatched_questions": unmatched_questions,
        }

    def _get_most_likely_answer(self, io_output_list: List[str]) -> Tuple[str, float]:
        """
        从多个候选答案中选择最可能的答案
        
        【核心选择逻辑】：
        使用评估器的投票机制从多个候选答案中选择最可能的答案。
        
        【判断依据】：
        1. 答案聚类：将相似的答案归为一组（使用check_answers_equiv判断等价性）
        2. 投票选择：选择出现频次最高的答案组
        3. 置信度计算：置信度 = 最高频次组的答案数量 / 总答案数量
        
        【具体步骤】：
        1. 从每个候选回答中提取最终答案
        2. 将等价答案归为同一组（例如"42"、"42.0"、"forty-two"可能等价）
        3. 统计每组答案的出现频次
        4. 选择频次最高的答案组作为最终答案
        5. 计算该组的置信度（频次/总数）
        
        【特殊情况处理】：
        - 只有一个候选答案：直接返回，置信度为1.0
        - 多个候选答案：通过投票机制选择最可能的答案
        - 无效答案：会被evaluator.find_most_confident_answer自动过滤
        
        Args:
            io_output_list (List[str]): 候选答案列表
            
        Returns:
            Tuple[str, float]: (最可能的答案, 置信度)
            
        示例：
            输入: ["42", "42.0", "43", "42"]
            处理: {"42": 3次, "43": 1次}  # "42"和"42.0"被认为等价
            输出: ("42", 0.75)  # 42出现3/4=75%的频次
        """
        assert len(io_output_list) > 0

        if len(io_output_list) == 1:
            # 只有一个候选答案，直接返回，置信度为1.0
            # 【文本选择】：单候选情况下，直接返回该候选的完整文本
            most_confident_answer_full_completion = io_output_list[0]
            confidence = 1
        else:
            # 多个候选答案，使用评估器的投票机制选择最可能的答案
            # find_most_confident_answer会：
            # 1. 提取每个候选的最终答案
            # 2. 将等价答案归为同一组
            # 3. 选择出现频次最高的答案组
            # 4. 返回该组的代表答案和置信度
            
            # 【关键】：这里返回的是完整的推理文本，而不仅仅是答案
            # most_confident_answer_full_completion 包含完整的推理过程
            # 例如："Let's solve step by step. First... Therefore, the answer is 42."
            # 而不仅仅是 "42"
            
            # 【文本选择逻辑】：
            # - 如果没有潜在分数权重：选择频次最高答案组的第一个完整文本
            # - 如果有潜在分数权重：选择加权分数最高的完整文本
            # 这确保了返回的文本不仅答案正确，而且推理过程具有代表性
            
            _, most_confident_answer_full_completion, _, confidence = self.evaluator.find_most_confident_answer(
                io_output_list
            )
            assert confidence > 0

        return most_confident_answer_full_completion, confidence

    def _fewshot_cot_answer_question(self, question: str, paraphrased: bool, num_return: int, hint: str = None):
        """
        使用少样本思维链方法回答问题
        
        【核心功能】：
        基于少样本学习(few-shot learning)和思维链(chain-of-thought)技术，
        为给定问题生成多个候选答案。
        
        【方法原理】：
        1. 模板构建：使用预定义的提示模板，包含示例和指令
        2. 上下文增强：可选择性地添加提示信息以改善生成质量
        3. 批量生成：一次性生成多个候选答案以提高覆盖率
        4. 停止控制：使用特定停止词确保生成内容的完整性
        
        【参数说明】：
        - question: 待回答的问题
        - paraphrased: 是否使用重述版本的提示模板
        - num_return: 生成候选答案的数量
        - hint: 可选的提示信息，用于引导生成方向
        
        Args:
            question (str): 待回答的问题
            paraphrased (bool): 是否使用重述版本的提示模板
            num_return (int): 生成候选答案的数量
            hint (str, optional): 提示信息，用于引导答案生成
            
        Returns:
            tuple: (输入提示, 清理后的输出列表)
                - io_input: 发送给模型的完整输入提示
                - cleaned_io_output_list: 去除首尾空白的答案列表
        """
        fewshot_cot_prompt = self.fewshot_cot_prompt if not paraphrased else self.fewshot_cot_prompt_rephrased
        question += "\n\n" + hint if hint is not None else ""
        io_input = self.fewshot_cot_config["prompt_template"].format(examples=fewshot_cot_prompt, instruction=question)
        io_output_list = self.io.generate(
            io_input,
            num_return=num_return,
            max_tokens=self.max_tokens,
            stop_tokens=self.fewshot_cot_config["stop_tokens"],
        )
        cleaned_io_output_list = [io_output.strip() for io_output in io_output_list]  #! cleaning
        return io_input, cleaned_io_output_list

    def generate_direct_answers(self, user_question: str, paraphrased: bool, hint: str):
        """
        生成直接答案节点
        
        使用少样本思维链（few-shot chain-of-thought）方法直接回答用户问题。
        这是推理过程中的终端动作，用于在当前状态下给出最终答案。
        
        【投票策略说明】：
        - 使用mcts_num_last_votes（默认32）而非num_votes（默认10）
        - 原因：直接答案是推理的最终结果，需要更高的可靠性
        - 更多候选答案通过投票机制能显著提高答案质量和一致性
        
        Args:
            user_question (str): 用户问题
            paraphrased (bool): 是否使用重述版本的提示
            hint (str): 提示信息，来自之前的推理步骤
            
        Returns:
            tuple: (直接答案列表, 置信度列表)
        """
        direct_answer_list, value_list = [], []

        #! few shot cot - 使用少样本思维链生成答案
        # 【关键】：使用mcts_num_last_votes确保终端答案的高质量
        num_return = self.mcts_num_last_votes
        io_input, cleaned_io_output_list = self._fewshot_cot_answer_question(
            question=user_question, paraphrased=paraphrased, num_return=num_return, hint=hint
        )

        try:
            # 从多个候选答案中选择最可能的答案
            most_likely_answer, likelihood = self._get_most_likely_answer(cleaned_io_output_list)
        except Exception as e:
            raise GeneratorError(
                source="generate direct answer from: few shot cot",
                io_input=io_input,
                io_output_list=cleaned_io_output_list,
            )

        direct_answer_list.append(most_likely_answer)
        value_list.append(likelihood)

        return direct_answer_list, value_list

    def generate_subquestions(
        self,
        user_question: str,
        solution_trace: Dict[int, Dict[str, str]],
        paraphrased: bool,
    ):
        """
        生成子问题节点（A3动作：问题分解）
        
        【核心功能】：
        将复杂的数学问题分解为更简单的子问题，并为每个子问题生成相应的答案。
        这是MCTS推理中的关键策略，通过分而治之的方式处理复杂问题。
        
        【执行流程】：
        1. 子问题生成：基于当前解决轨迹生成下一个逻辑子问题
        2. 子答案生成：为每个生成的子问题提供候选答案
        3. 投票选择：通过多候选投票机制选择最可能的答案
        4. 潜在评估：可选地评估每个子问题对最终答案的贡献度
        
        【投票策略】：
        根据子问题是否为终端问题采用不同的投票数量：
        - 终端子问题：使用mcts_num_last_votes（更多候选，确保质量）
        - 非终端子问题：使用num_votes（适中候选，平衡效率）
        
        【潜在分数评估】：
        当启用时，会评估每个子问题-答案对对最终解决方案的潜在贡献，
        用于后续的节点选择和搜索引导。
        
        Args:
            user_question (str): 原始用户问题
            solution_trace (Dict[int, Dict[str, str]]): 当前解决轨迹
            paraphrased (bool): 是否使用重述版本的提示
            
        Returns:
            tuple: (子问题列表, 子答案列表, 置信度列表, 潜在答案列表)
                - subquestion_list: 生成的子问题列表
                - subanswer_list: 对应的子答案列表
                - value_list: 每个子答案的置信度
                - potential_answers_list: 潜在答案评估结果（如果启用）
        """
        subquestion_list, subanswer_list, value_list = [], [], []
        decompose_prompt = self.decompose_prompt if not paraphrased else self.decompose_prompt_rephrased

        #! 第一步：生成子问题
        # 构建包含已有子问题和答案的上下文
        existing_subquestions_and_subanswers, next_subquestion_id = concat_subqs_and_subas(
            solution_trace, self.question_index
        )
        # 构建子问题生成的输入提示
        io_input = (
            decompose_prompt
            + "\n\n"
            + f"Question {self.question_index}: {user_question}"
            + "\n"
            + existing_subquestions_and_subanswers
            + f"Question {self.question_index}.{next_subquestion_id}:"
        )
        # 生成多个候选子问题
        io_output_list = self.io.generate(
            io_input,
            max_tokens=128,
            num_return=self.num_subquestions,
            stop_tokens=[
                "\n",
                "\n\n",
                "Answer",
                "Answer ",
                f"Answer {self.question_index}.{next_subquestion_id}",
                f"Answer {self.question_index}.{next_subquestion_id}:",
                f"Answer {self.question_index}.{next_subquestion_id}: ",
            ],
        )

        # 清理生成的子问题文本
        # subquestion_list = [io_output.split("?")[0] + "?" for io_output in io_output_list]  # cleaning, you might wanna modify this
        subquestion_list = [o.strip() for o in io_output_list]

        #! 第二步：为生成的子问题生成答案
        # 为每个子问题构建答案生成的输入提示
        io_input_list = []
        num_return_list = []  # 为每个子问题存储对应的候选数量
        
        for subquestion in subquestion_list:
            io_input = (
                decompose_prompt
                + "\n\n"
                + f"Question {self.question_index}: {user_question}"
                + "\n"
                + existing_subquestions_and_subanswers
                + f"Question {self.question_index}.{next_subquestion_id}: "
                + subquestion
                + "\n"
                + f"Answer {self.question_index}.{next_subquestion_id}:"
            )
            io_input_list.append(io_input)
            
            # 【核心投票逻辑】：根据子问题类型选择不同的投票策略
            if reach_terminal_subquestion(subquestion=subquestion, user_question=user_question):
                # 终端子问题：使用更多候选确保最终答案质量
                # 终端子问题通常包含"Now we can answer"或回到原问题核心
                num_return_list.append(self.mcts_num_last_votes)  # 默认32个候选
            else:
                # 非终端子问题：使用较少候选平衡效率和质量
                # 中间步骤的答案主要用于推理过程的延续
                num_return_list.append(self.num_votes)  # 默认10个候选

        # 批量生成子问题的答案，为每个子问题使用对应的候选数量
        io_output_list = []
        for i, (io_input, num_return) in enumerate(zip(io_input_list, num_return_list)):
            outputs = self.io.generate(
                io_input,
                max_tokens=512,
                num_return=num_return,
                stop_tokens=[
                    "\n",
                    "\n\n",
                    f"Question {self.question_index}.{next_subquestion_id + 1}",
                ],
            )
            io_output_list.append(outputs)

        # 清理生成的答案文本
        cleaned_io_output_list = [
            [io_output.strip() for io_output in io_output_group] for io_output_group in io_output_list
        ]

        # 第三步：从候选答案中选择最可能的答案
        for i, cleaned_io_output_group in enumerate(cleaned_io_output_list):
            try:
                most_likely_answer, likelihood = self._get_most_likely_answer(cleaned_io_output_group)
            except Exception as e:
                raise GeneratorError(
                    source="generate answer to subquestions",
                    io_input=io_input_list[i],
                    io_output_list=cleaned_io_output_group,
                )
            subanswer_list.append(most_likely_answer)
            value_list.append(likelihood)

        # 确保三个列表长度一致
        assert len(subquestion_list) == len(subanswer_list) == len(value_list)

        #! 第四步：生成潜在分数评估（如果启用）
        potential_answers_list: List[List[str]] = []
        if self.enable_potential_score:
            # 为每个子问题-答案对评估其对最终答案的潜在贡献
            for subq, suba in zip(subquestion_list, subanswer_list):
                if reach_terminal_subquestion(subq, user_question):
                    # 终端子问题已经是最终答案，无需潜在分数评估
                    potential_answers_list.append(None)
                else:
                    # 构建包含新子问题和答案的响应前缀
                    response_prefix = make_response_prefix(
                        solution_trace, Node_Type.SUBQUESTION, new_subq=subq, new_suba=suba
                    )
                    potential_score_input = "Question: " + user_question + "\nAnswer: " + response_prefix

                    # 【潜在分数评估投票策略】：统一使用num_votes
                    # 原因：潜在分数评估是中间过程，用于节点选择而非最终答案
                    # 目的：快速评估节点潜在价值，不需要最终答案级别的精确性
                    potential_score_output = self.io.generate(
                        potential_score_input,
                        num_return=self.num_votes,
                        max_tokens=128,
                        stop_tokens=self.fewshot_cot_config["stop_tokens"],
                    )
                    # 构建完整的潜在答案评估输入
                    potential_score_input2 = [
                        "Question: "
                        + user_question
                        + "\nAnswer: "
                        + response_prefix
                        + z
                        + "\nTherefore, the answer (arabic numerals) is"
                        for z in potential_score_output
                    ]
                    # 生成潜在答案
                    cleaned_io_output_list = self.io.generate(
                        potential_score_input2,
                        num_return=1,
                        max_tokens=128,
                        stop_tokens=self.fewshot_cot_config["stop_tokens"],
                    )
                    cleaned_io_output_list = [z[0] for z in cleaned_io_output_list]

                    # 提取潜在答案并添加到列表
                    potential_answers_list.append(
                        [self.evaluator.extract_answer_from_model_completion(o) for o in cleaned_io_output_list]
                    )
        else:
            # 未启用潜在分数评估，为每个子问题设置None
            potential_answers_list = [None] * len(subquestion_list)

        return subquestion_list, subanswer_list, value_list, potential_answers_list

    def generate_re_subanswers(
        self,
        user_question: str,
        solution_trace: Dict[int, Dict[str, str]],
        paraphrased: bool,
    ):
        """
        生成重新回答子问题节点（A4动作：重新回答）
        
        【核心功能】：
        对已有的子问题提供新的答案，用于探索不同的解决路径。
        这是MCTS推理中的多样化策略，通过重新回答增加解决方案的覆盖率。
        
        【应用场景】：
        1. 当前子答案质量不高时，尝试生成更好的答案
        2. 探索同一子问题的不同解决角度
        3. 在MCTS搜索中提供更多的分支选择
        
        【执行流程】：
        1. 定位最后一个子问题：从解决轨迹中获取最近的子问题
        2. 重新生成答案：使用少样本思维链方法重新回答
        3. 投票选择：通过多候选投票选择最可能的新答案
        4. 潜在评估：评估新答案对最终解决方案的潜在贡献
        
        【投票策略】：
        使用num_votes进行投票选择，平衡答案质量和生成效率。
        
        Args:
            user_question (str): 原始用户问题
            solution_trace (Dict[int, Dict[str, str]]): 当前解决轨迹
            paraphrased (bool): 是否使用重述版本的提示
            
        Returns:
            tuple: (重新回答列表, 置信度列表, 潜在答案列表)
                - re_subanswer_list: 重新生成的子答案列表
                - value_list: 每个重新回答的置信度
                - potential_answers_list: 潜在答案评估结果（如果启用）
        """
        re_subanswer_list, value_list = [], []

        user_question_context, _ = split_user_question(user_question)

        last_subquestion_id = int(sorted(solution_trace.keys())[-1])
        last_subquestion = solution_trace[last_subquestion_id]["subquestion"]

        #! few shot cot
        question = (
            f"{user_question_context} {last_subquestion}"
            if not paraphrased
            else f"{user_question_context} Question: {last_subquestion}"
        )
        # 【投票策略】：重新回答子问题使用num_votes
        # 原因：这是对已有子问题的重新回答，属于中间推理步骤
        # 目的：提供多样化的答案选择，而不需要最终答案级别的高可靠性
        io_input, cleaned_io_output_list = self._fewshot_cot_answer_question(
            question=question, paraphrased=paraphrased, num_return=self.num_votes
        )
        try:
            most_likely_answer, likelihood = self._get_most_likely_answer(cleaned_io_output_list)
        except Exception as e:
            raise GeneratorError(
                source="generate re-subanswers: few shot cot",
                io_input=io_input,
                io_output_list=cleaned_io_output_list,
            )
        re_subanswer_list.append(most_likely_answer)
        value_list.append(likelihood)

        #! generate potential answer to the user question
        potential_answers_list: List[List[str]] = []
        if self.enable_potential_score:
            solution_trace_copy = deepcopy(solution_trace)
            for re_suba in re_subanswer_list:
                solution_trace_copy[last_subquestion_id]["subanswer"] = {"text": re_suba}
                response_prefix = make_response_prefix(solution_trace_copy, Node_Type.SUBQUESTION)
                potential_score_input = "Question: " + user_question + "\nAnswer: " + response_prefix

                potential_score_output = self.io.generate(
                    potential_score_input,
                    num_return=self.num_votes,
                    max_tokens=128,
                    stop_tokens=self.fewshot_cot_config["stop_tokens"],
                )
                potential_score_input2 = [
                    "Question: "
                    + user_question
                    + "\nAnswer: "
                    + response_prefix
                    + z
                    + "\nTherefore, the answer (arabic numerals) is"
                    for z in potential_score_output
                ]
                cleaned_io_output_list = self.io.generate(
                    potential_score_input2,
                    num_return=1,
                    max_tokens=128,
                    stop_tokens=self.fewshot_cot_config["stop_tokens"],
                )
                cleaned_io_output_list = [z[0] for z in cleaned_io_output_list]

                potential_answers_list.append(
                    [self.evaluator.extract_answer_from_model_completion(o) for o in cleaned_io_output_list]
                )
        else:
            potential_answers_list = [None] * len(re_subanswer_list)

        return re_subanswer_list, value_list, potential_answers_list

    def generate_rephrased_user_question(self, user_question: str):
        """
        生成重述用户问题节点（A5动作：问题重述）
        
        【核心功能】：
        将原始用户问题重新表述为不同的形式，以获得新的理解角度。
        这是MCTS推理中的视角转换策略，通过改变问题表述来启发新的解决思路。
        
        【重述策略】：
        使用预定义的重述模板将原问题转换为条件列表形式，
        例如："求解x+y=10, x-y=2中的x"转换为"给定条件列表，请回答问题。条件1: x+y=10..."
        
        【应用价值】：
        1. 视角转换：从不同角度理解同一问题
        2. 提示优化：重述后的问题可能更适合模型理解
        3. 搜索多样性：为MCTS提供不同的搜索起点
        
        【执行流程】：
        1. 问题重述：使用重述模板转换原问题
        2. 潜在评估：评估重述问题对解决方案的潜在贡献
        
        Args:
            user_question (str): 原始用户问题
            
        Returns:
            tuple: (重述问题列表, 潜在答案列表)
                - rephrased_user_question_list: 重述后的问题列表
                - potential_answers_list: 潜在答案评估结果（如果启用）
        """
        rephrased_user_question_list = []
        io_input = self.rephrasing_prompt_template
        io_input += "\n\n"
        io_input += "Original Question: " + user_question + "\n"
        io_input += "Rephrased Question: Given a list of conditions, please answer the question. Condition 1: "
        io_output = self.io.generate(model_input=io_input, max_tokens=512, num_return=1, stop_tokens=["\n", "\n\n"])[0]
        io_output = "Given a list of conditions, please answer the question. Condition 1: " + io_output
        rephrased_user_question_list.append(io_output)

        #! generate potential answer to the user question
        potential_answers_list: List[List[str]] = []  # essentially direct answer list
        if self.enable_potential_score:
            response_prefix = make_response_prefix(None, None)
            potential_score_input = "Question: " + rephrased_user_question_list[0] + "\nAnswer: " + response_prefix
            # 【潜在分数评估投票策略】：统一使用num_votes
            # 原因：潜在分数评估是中间过程，用于节点选择而非最终答案
            # 目的：快速评估节点潜在价值，不需要最终答案级别的精确性
            potential_score_output = self.io.generate(
                potential_score_input,
                num_return=self.num_votes,
                max_tokens=128,
                stop_tokens=self.fewshot_cot_config["stop_tokens"],
            )
            potential_score_input2 = [
                "Question: "
                + rephrased_user_question_list[0]
                + "\nAnswer: "
                + response_prefix
                + z
                + "\nTherefore, the answer (arabic numerals) is"
                for z in potential_score_output
            ]
            cleaned_io_output_list = self.io.generate(
                potential_score_input2, num_return=1, max_tokens=128, stop_tokens=self.fewshot_cot_config["stop_tokens"]
            )
            cleaned_io_output_list = [z[0] for z in cleaned_io_output_list]

            potential_answers_list.append(
                [self.evaluator.extract_answer_from_model_completion(o) for o in cleaned_io_output_list]
            )
        else:
            potential_answers_list = [None] * len(rephrased_user_question_list)

        return rephrased_user_question_list, potential_answers_list

    def generate_ost_step(
        self,
        user_question: str,
        solution_trace: Dict[int, Dict[str, str]],
        paraphrased: bool,
        parent_is_subquestion: bool,
    ):
        """
        生成一步思考节点（A1动作：单步推理）
        
        【核心功能】：
        生成单步推理思考，逐步推进问题解决过程。
        这是MCTS推理中的渐进式策略，通过逐步思考来构建完整的解决方案。
        
        【思考类型】：
        1. 直接思考：从用户问题开始的逐步推理
        2. 子问题思考：基于已有子问题的进一步推理
        
        【执行流程】：
        1. 上下文构建：根据父节点类型构建不同的思考上下文
           - 子问题父节点：将子问题和答案转换为思考步骤
           - 其他父节点：直接使用现有的思考步骤
        2. 步骤生成：使用一步思考模板生成下一步推理
        3. 潜在评估：评估每个思考步骤对最终解决方案的贡献
        
        【应用场景】：
        - 复杂问题的逐步分析
        - 推理过程的细化和深入
        - 为后续直接答案生成提供更好的上下文
        
        Args:
            user_question (str): 原始用户问题
            solution_trace (Dict[int, Dict[str, str]]): 当前解决轨迹
            paraphrased (bool): 是否使用重述版本的提示
            parent_is_subquestion (bool): 父节点是否为子问题节点
            
        Returns:
            tuple: (思考步骤列表, 潜在答案列表)
                - ost_step_list: 生成的一步思考列表
                - potential_answers_list: 潜在答案评估结果（如果启用）
        """
        ost_step_list = []
        if parent_is_subquestion:
            existing_ost_steps, next_ost_step_id = concat_subqs_subas_as_ost_steps(solution_trace)
        else:
            existing_ost_steps, next_ost_step_id = concat_ost_steps(solution_trace)
        io_input = (
            self.fewshot_ost_config["prompt_template"].format(
                examples=self.fewshot_ost_prompt if not paraphrased else self.fewshot_ost_prompt_rephrased,
                instruction=user_question,
            )
            + existing_ost_steps
            + f"Step {next_ost_step_id}:"
        )
        io_output_list = self.io.generate(
            model_input=io_input, max_tokens=256, num_return=self.num_a1_steps, stop_tokens=["\n", "\n\n"]
        )
        ost_step_list = [io_output.strip() for io_output in io_output_list]

        #! generate potential answer to the user question
        potential_answers_list: List[List[str]] = []  # essentially direct answer list
        if self.enable_potential_score:
            for ost_step in ost_step_list:
                response_prefix = make_response_prefix(solution_trace, Node_Type.OST_STEP, new_ost_step=ost_step)

                potential_score_input = "Question: " + user_question + "\nAnswer: " + response_prefix

                # 【潜在分数评估投票策略】：统一使用num_votes
                # 原因：潜在分数评估是中间过程，用于节点选择而非最终答案
                # 目的：快速评估节点潜在价值，不需要最终答案级别的精确性
                potential_score_output = self.io.generate(
                    potential_score_input,
                    num_return=self.num_votes,
                    max_tokens=128,
                    stop_tokens=self.fewshot_cot_config["stop_tokens"],
                )
                potential_score_input2 = [
                    "Question: "
                    + user_question
                    + "\nAnswer: "
                    + response_prefix
                    + z
                    + "\nTherefore, the answer (arabic numerals) is"
                    for z in potential_score_output
                ]
                cleaned_io_output_list = self.io.generate(
                    potential_score_input2,
                    num_return=1,
                    max_tokens=128,
                    stop_tokens=self.fewshot_cot_config["stop_tokens"],
                )
                cleaned_io_output_list = [z[0] for z in cleaned_io_output_list]

                potential_answers_list.append(
                    [self.evaluator.extract_answer_from_model_completion(o) for o in cleaned_io_output_list]
                )
        else:
            potential_answers_list = [None] * len(ost_step_list)

        return ost_step_list, potential_answers_list


class Reasoning_MCTS_Node(MCTS_Node):
    """
    推理任务专用的MCTS节点类
    
    继承自MCTS_Node基类，专门用于数学推理任务。
    支持多种节点类型和推理动作，维护完整的推理轨迹。
    
    主要特性：
    - 支持6种节点类型：用户问题、重述问题、直接答案、子问题、重新回答、一步思考
    - 维护解决方案轨迹，记录从根节点到当前节点的完整推理过程
    - 支持潜在分数评估，提高解决方案选择的准确性
    - 实现推理动作的生成和管理
    
    核心方法：
    - find_children: 生成子节点（实现各种推理动作）
    - is_terminal: 判断是否为终端节点
    - calculate_reward: 计算节点奖励值
    - is_valid_solution_node: 判断是否为有效解决方案节点
    """
    
    def __init__(
        self,
        parent: "Reasoning_MCTS_Node",
        depth: int,
        node_type: Node_Type,
        verbose: bool = False,
        # --- For instantiating root node ---
        node_value: float = None,
        generator: Generator = None,
        disable_a5: bool = None,
        user_question: str = None,
        max_depth_allowed: int = None,
        disable_a1: bool = None,
        # -----------------------------------
        # --- For instantiating REPHRASED_USER_QUESTION node ---
        rephrased_user_question: str = None,
        # ------------------------------------------------------
        expected_answer: str = None,
        # --- For instantiating DIRECT_ANSWER node ---
        direct_answer: str = None,
        # --------------------------------------------
        # --- For instantiating SUBQUESTION node ---
        subquestion: str = None,
        subanswer: str = None,
        is_new_subquestion: bool = None,
        # ------------------------------------------
        # --- For instantiating RE_SUBANSWER node ---
        re_subanswer: str = None,
        # -------------------------------------------
        # --- For instantiating OST_STEP node ---
        ost_step: str = None,
        # ---------------------------------------
        # --- For node selection (not in sanity checks yet) ---
        enable_potential_score: bool = None,
        potential_answers: List[str] = None,
    ) -> None:
        """
        初始化推理MCTS节点
        
        根据节点类型初始化不同的属性，并维护推理轨迹和状态。
        
        Args:
            parent: 父节点
            depth: 节点深度
            node_type: 节点类型
            verbose: 是否输出详细信息
            
            # 根节点专用参数
            node_value: 节点价值（终端节点的置信度）
            generator: 推理生成器
            disable_a5: 是否禁用A5动作（问题重述）
            user_question: 用户问题
            max_depth_allowed: 最大深度限制
            disable_a1: 是否禁用A1动作（一步思考）
            
            # 各种节点类型专用参数
            rephrased_user_question: 重述的用户问题
            expected_answer: 预期答案（用于评估）
            direct_answer: 直接答案内容
            subquestion: 子问题内容
            subanswer: 子答案内容
            is_new_subquestion: 是否为新的子问题
            re_subanswer: 重新回答的子答案
            ost_step: 一步思考的内容
            
            # 潜在分数评估相关
            enable_potential_score: 是否启用潜在分数评估
            potential_answers: 潜在答案列表
        """
        super().__init__()

        #! sanity checks
        try:
            assert depth is not None
            assert node_type is not None
            if node_value is not None:
                assert node_value > 0, breakpoint()

            if node_type is Node_Type.USER_QUESTION:
                assert depth == 0
                assert all(
                    attr is None
                    for attr in [
                        parent,
                        node_value,
                        rephrased_user_question,
                        direct_answer,
                        subquestion,
                        subanswer,
                        is_new_subquestion,
                        re_subanswer,
                        ost_step,
                    ]
                )
                assert all(
                    attr is not None
                    for attr in [generator, disable_a5, user_question, expected_answer, max_depth_allowed, disable_a1]
                )
            elif node_type is Node_Type.REPHRASED_USER_QUESTION:
                assert depth == 1
                assert all(
                    attr is None
                    for attr in [
                        node_value,
                        generator,
                        disable_a5,
                        user_question,
                        expected_answer,
                        direct_answer,
                        subquestion,
                        subanswer,
                        is_new_subquestion,
                        re_subanswer,
                        ost_step,
                        max_depth_allowed,
                        disable_a1,
                    ]
                )
                assert all(attr is not None for attr in [parent, rephrased_user_question])
            elif node_type is Node_Type.DIRECT_ANSWER:
                assert depth > 0
                assert all(
                    attr is None
                    for attr in [
                        generator,
                        disable_a5,
                        user_question,
                        expected_answer,
                        subquestion,
                        subanswer,
                        is_new_subquestion,
                        re_subanswer,
                        ost_step,
                        max_depth_allowed,
                        disable_a1,
                    ]
                )
                assert all(attr is not None for attr in [parent, node_value, direct_answer])
            elif node_type is Node_Type.SUBQUESTION:
                assert depth > 0
                assert all(
                    attr is None
                    for attr in [
                        generator,
                        disable_a5,
                        user_question,
                        expected_answer,
                        direct_answer,
                        re_subanswer,
                        ost_step,
                        max_depth_allowed,
                        disable_a1,
                    ]
                )
                assert all(
                    attr is not None for attr in [parent, node_value, subquestion, subanswer, is_new_subquestion]
                )
            elif node_type is Node_Type.RE_SUBANSWER:
                assert depth > 0
                assert all(
                    attr is None
                    for attr in [
                        generator,
                        disable_a5,
                        user_question,
                        expected_answer,
                        direct_answer,
                        subquestion,
                        subanswer,
                        is_new_subquestion,
                        ost_step,
                        max_depth_allowed,
                        disable_a1,
                    ]
                )
                assert all(attr is not None for attr in [parent, node_value, re_subanswer])
            elif node_type is Node_Type.OST_STEP:
                assert depth > 0
                assert all(
                    attr is None
                    for attr in [
                        node_value,
                        generator,
                        disable_a5,
                        user_question,
                        rephrased_user_question,
                        expected_answer,
                        direct_answer,
                        subquestion,
                        subanswer,
                        is_new_subquestion,
                        re_subanswer,
                        max_depth_allowed,
                        disable_a1,
                    ]
                )
                assert all(attr is not None for attr in [parent, ost_step])
        except AssertionError:
            print(f"Instantiating node with type {node_type} failed!")
            breakpoint()
            exit()

        #! attributes
        self.parent = parent  # if parent is None, then the node is the root
        self.children: List["Reasoning_MCTS_Node"] = []
        self.depth = depth
        self.node_type = node_type
        self.node_value = node_value
        self.direct_answer = direct_answer
        self.subquestion = subquestion
        self.subanswer = subanswer
        self.is_new_subquestion = is_new_subquestion
        self.re_subanswer = re_subanswer
        self.ost_step = ost_step

        if parent is None:  # root
            self.verbose = verbose
            self.user_question = user_question
            self.expected_answer = expected_answer
            self.generator = generator
            self.disable_a5 = disable_a5
            self.question_index = generator.question_index
            self.max_depth_allowed = max_depth_allowed
            self.disable_a1 = disable_a1
            self.enable_potential_score = enable_potential_score
        else:  # inherit from parent
            self.verbose = parent.verbose
            self.user_question = parent.user_question
            self.expected_answer = parent.expected_answer
            self.generator = parent.generator
            self.disable_a5 = parent.disable_a5
            self.question_index = parent.generator.question_index
            self.max_depth_allowed = parent.max_depth_allowed
            self.disable_a1 = parent.disable_a1
            self.enable_potential_score = parent.enable_potential_score

        #! keep track of paraphrasing
        if node_type is Node_Type.USER_QUESTION:
            self.paraphrased = False
        elif node_type is Node_Type.REPHRASED_USER_QUESTION:
            self.paraphrased = True
            self.user_question = rephrased_user_question
        else:
            assert parent is not None
            self.paraphrased = parent.paraphrased

        #! record number of subquestions till now
        if parent is None:  # root
            self.subquestion_counter = 0
        else:
            if node_type is Node_Type.SUBQUESTION and is_new_subquestion:
                self.subquestion_counter = parent.subquestion_counter + 1
            else:
                self.subquestion_counter = parent.subquestion_counter

        #! record number of one-step thought steps till now
        if parent is None:  # root
            self.ost_step_counter = 0
        else:
            if node_type is Node_Type.OST_STEP:
                self.ost_step_counter = parent.ost_step_counter + 1
            else:
                self.ost_step_counter = parent.ost_step_counter

        #! record solution trace from root to the current node. key: subquestion id
        if parent is None:  # root
            assert self.node_type is Node_Type.USER_QUESTION
            self.solution_trace: Dict[int, Dict[str, str]] = {0: {"user_question": user_question, "ost_step": {}}}
        else:
            assert self.node_type is not Node_Type.USER_QUESTION
            self.solution_trace = deepcopy(parent.solution_trace)

            if node_type is Node_Type.REPHRASED_USER_QUESTION:
                self.solution_trace[0]["user_question"] = rephrased_user_question
            elif node_type is Node_Type.DIRECT_ANSWER:
                assert self.subquestion_counter in self.solution_trace.keys()
                assert self.subquestion_counter == parent.subquestion_counter
                self.solution_trace[self.subquestion_counter]["direct_answer"] = {
                    "text": direct_answer,
                    "value": node_value,
                }
            elif node_type is Node_Type.SUBQUESTION:
                assert is_new_subquestion and self.subquestion_counter == parent.subquestion_counter + 1
                self.solution_trace[self.subquestion_counter] = {
                    "subquestion": subquestion,
                    "subanswer": {"text": subanswer, "value": node_value},
                    "ost_step": {},
                }
            elif node_type is Node_Type.RE_SUBANSWER:
                assert parent.subquestion is not None
                assert self.subquestion_counter == parent.subquestion_counter
                assert self.solution_trace[self.subquestion_counter]["subquestion"] == parent.subquestion
                self.solution_trace[self.subquestion_counter]["subanswer"] = {"text": re_subanswer, "value": node_value}
            elif node_type is Node_Type.OST_STEP:
                assert "ost_step" in self.solution_trace[self.subquestion_counter].keys()
                self.solution_trace[self.subquestion_counter]["ost_step"][self.ost_step_counter] = ost_step

        #! potential_score for intermediate nodes (only used for node selection)
        if self.enable_potential_score:
            self.potential_answers = potential_answers
            self.potential_score = 0
            if parent is None:  # root
                assert self.node_type is Node_Type.USER_QUESTION
                self.potential_answers_history = {}
            else:
                assert self.node_type is not Node_Type.USER_QUESTION
                self.potential_answers_history = deepcopy(parent.potential_answers_history)
                self.potential_answers_history[self.depth] = potential_answers

    def __str__(self) -> str:
        """
        节点的字符串表示
        
        【功能说明】：
        为节点提供简洁的字符串表示，用于调试和日志输出。
        
        【格式规则】：
        使用"节点类型缩写-节点ID"的格式，例如：
        - U-0: 用户问题节点，ID为0
        - SQ-5: 子问题节点，ID为5
        - DA-12: 直接答案节点，ID为12
        
        【节点类型缩写】：
        - U: USER_QUESTION (用户问题)
        - RU: REPHRASED_USER_QUESTION (重述用户问题)
        - DA: DIRECT_ANSWER (直接答案)
        - SQ: SUBQUESTION (子问题)
        - RS: RE_SUBANSWER (重新回答)
        - TS: OST_STEP (一步思考)
        
        Returns:
            str: 节点的字符串表示
        """
        type2str = {
            Node_Type.USER_QUESTION: "U",
            Node_Type.REPHRASED_USER_QUESTION: "RU",
            Node_Type.DIRECT_ANSWER: "DA",
            Node_Type.SUBQUESTION: "SQ",
            Node_Type.RE_SUBANSWER: "RS",
            Node_Type.OST_STEP: "TS",
        }
        return f"{type2str[self.node_type]}-{self.id}"

    def _create_children(self):
        """
        创建子节点的核心方法
        
        【功能说明】：
        根据当前节点类型，执行相应的推理动作（A1-A5）来生成子节点。
        这是MCTS扩展阶段的核心实现，为每种节点类型定义了可执行的动作集合。
        
        【推理动作映射】：
        - A1: 生成一步思考节点 (OST_STEP)
        - A2: 生成直接答案节点 (DIRECT_ANSWER)
        - A3: 生成子问题节点 (SUBQUESTION)
        - A4: 生成重新回答节点 (RE_SUBANSWER)
        - A5: 生成重述问题节点 (REPHRASED_USER_QUESTION)
        
        【节点类型与动作关系】：
        - USER_QUESTION: 支持A1, A2, A3, A5
        - REPHRASED_USER_QUESTION: 支持A1, A2, A3
        - SUBQUESTION: 支持A1, A2, A3, A4
        - RE_SUBANSWER: 支持A1, A2, A3
        - OST_STEP: 支持A1, A2
        - DIRECT_ANSWER: 终端节点，无子节点
        
        【内部方法说明】：
        包含多个do_action_*内部方法，每个方法负责执行特定的推理动作：
        - do_action_generate_direct_answers: 执行A2动作
        - do_action_generate_subquestions: 执行A3动作
        - do_action_generate_re_subanswers: 执行A4动作
        - do_action_generate_rephrased_user_question: 执行A5动作
        - do_action_generate_ost_step: 执行A1动作
        
        Returns:
            List[Reasoning_MCTS_Node]: 生成的子节点列表
        """
        def do_action_generate_direct_answers():
            """
            执行A2动作：生成直接答案节点
            
            【功能说明】：
            基于当前推理状态直接生成问题的最终答案。
            这是推理过程的终端动作，用于在当前上下文下给出解决方案。
            
            【提示策略】：
            - 根节点和重述问题节点：无额外提示，直接回答
            - 其他节点：使用make_hint()生成上下文提示，包含已有的推理轨迹
            
            【节点创建】：
            为每个生成的直接答案创建DIRECT_ANSWER类型的子节点，
            节点价值为答案的置信度。
            """
            verbose_print(f"---- Generating direct answers for node {self.id}...", self.verbose)

            #! ACTION: generate direct answer for the user question (w/ or w/o hint)
            if (
                self.node_type is not Node_Type.USER_QUESTION
                and self.node_type is not Node_Type.REPHRASED_USER_QUESTION
            ):
                hint = make_hint(self.solution_trace, self.node_type)
            else:
                hint = None

            (direct_answer_list, value_list) = self.generator.generate_direct_answers(
                user_question=self.user_question, paraphrased=self.paraphrased, hint=hint
            )
            for direct_answer, value in zip(direct_answer_list, value_list):
                if np.isnan(value) or value <= 0:
                    breakpoint()
                self.children.append(
                    Reasoning_MCTS_Node(
                        parent=self,
                        depth=self.depth + 1,
                        node_type=Node_Type.DIRECT_ANSWER,
                        node_value=value,
                        direct_answer=direct_answer,
                    )
                )

        def do_action_generate_subquestions():
            """
            执行A3动作：生成子问题节点
            
            【功能说明】：
            将复杂问题分解为更简单的子问题，并为每个子问题生成答案。
            这是问题分解策略的核心实现，通过分而治之处理复杂推理。
            
            【执行过程】：
            1. 基于当前解决轨迹生成新的子问题
            2. 为每个子问题生成对应的答案
            3. 计算潜在分数（如果启用）
            4. 创建SUBQUESTION类型的子节点
            
            【节点创建】：
            为每个子问题-答案对创建子节点，包含：
            - 子问题内容和答案
            - 答案置信度作为节点价值
            - 潜在答案评估结果（用于节点选择）
            
            【异常处理】：
            对于置信度异常的答案，设置最小值0.01以保证节点可用。
            """
            verbose_print(f"---- Generating subquestions for node {self.id}...", self.verbose)

            #! ACTION: generate new subquestions
            (subquestion_list, subanswer_list, value_list, potential_answers_list) = (
                self.generator.generate_subquestions(
                    user_question=self.user_question, solution_trace=self.solution_trace, paraphrased=self.paraphrased
                )
            )
            for subquestion, subanswer, value, potential_answers in zip(
                subquestion_list, subanswer_list, value_list, potential_answers_list
            ):
                if np.isnan(value) or value <= 0:
                    value = 0.01
                    # breakpoint()
                self.children.append(
                    Reasoning_MCTS_Node(
                        parent=self,
                        depth=self.depth + 1,
                        node_type=Node_Type.SUBQUESTION,
                        node_value=value,
                        subquestion=subquestion,
                        subanswer=subanswer,
                        is_new_subquestion=True,
                        potential_answers=deepcopy(potential_answers),
                    )
                )

        def do_action_generate_re_subanswers():
            """
            执行A4动作：生成重新回答节点
            
            【功能说明】：
            对最近的子问题重新生成答案，探索不同的解决路径。
            这是多样化策略的实现，通过重新回答增加解决方案的覆盖率。
            
            【适用场景】：
            - 当前子答案质量不理想时
            - 需要探索同一子问题的不同解决角度
            - 为MCTS提供更多搜索分支
            
            【执行过程】：
            1. 定位最后一个子问题
            2. 重新生成该子问题的答案
            3. 计算潜在分数（如果启用）
            4. 创建RE_SUBANSWER类型的子节点
            
            【节点创建】：
            为每个重新生成的答案创建子节点，包含：
            - 新的子答案内容
            - 答案置信度作为节点价值
            - 潜在答案评估结果
            """
            verbose_print(f"---- Generating re-subanswers for node {self.id}...", self.verbose)

            #! ACTION: re-generate subanswers for the previous subquestion
            (re_subanswer_list, value_list, potential_answers_list) = self.generator.generate_re_subanswers(
                user_question=self.user_question,
                solution_trace=self.solution_trace,
                paraphrased=self.paraphrased,
            )
            for re_subanswer, value, potential_answers in zip(re_subanswer_list, value_list, potential_answers_list):
                if np.isnan(value) or value <= 0:
                    breakpoint()
                self.children.append(
                    Reasoning_MCTS_Node(
                        parent=self,
                        depth=self.depth + 1,
                        node_type=Node_Type.RE_SUBANSWER,
                        node_value=value,
                        re_subanswer=re_subanswer,
                        potential_answers=deepcopy(potential_answers),
                    )
                )

        def do_action_generate_rephrased_user_question():
            """
            执行A5动作：生成重述问题节点
            
            【功能说明】：
            重新表述原始用户问题，从不同角度理解和解决问题。
            这是视角转换策略的实现，通过改变问题表述启发新的解决思路。
            
            【重述价值】：
            - 提供问题的不同理解角度
            - 可能更适合模型的理解和处理
            - 为MCTS搜索提供多样化的起点
            
            【执行过程】：
            1. 使用重述模板转换原问题
            2. 计算潜在分数（如果启用）
            3. 创建REPHRASED_USER_QUESTION类型的子节点
            
            【节点创建】：
            为每个重述问题创建子节点，包含：
            - 重述后的问题内容
            - 潜在答案评估结果（用于节点选择）
            
            【注意事项】：
            重述问题节点没有node_value，因为它不是终端节点。
            """
            verbose_print(f"---- Generating rephrased user question for node {self.id}...", self.verbose)

            #! ACTION: generate paraphrased question for the root question
            rephrased_user_question_list, potential_answers_list = self.generator.generate_rephrased_user_question(
                user_question=self.user_question
            )
            for rephrased_user_question, potential_answers in zip(rephrased_user_question_list, potential_answers_list):
                self.children.append(
                    Reasoning_MCTS_Node(
                        parent=self,
                        depth=self.depth + 1,
                        node_type=Node_Type.REPHRASED_USER_QUESTION,
                        rephrased_user_question=rephrased_user_question,
                        potential_answers=deepcopy(potential_answers),
                    )
                )

        def do_action_generate_ost_step(parent_is_subquestion=False):
            """
            执行A1动作：生成一步思考节点
            
            【功能说明】：
            生成单步推理思考，逐步推进问题解决过程。
            这是渐进式推理策略的实现，通过逐步思考构建完整解决方案。
            
            【思考类型】：
            - 直接思考：从当前状态继续推理
            - 子问题思考：基于子问题的进一步分析
            
            【执行过程】：
            1. 根据parent_is_subquestion确定思考上下文
            2. 生成下一步推理思考
            3. 计算潜在分数（如果启用）
            4. 创建OST_STEP类型的子节点
            
            【节点创建】：
            为每个思考步骤创建子节点，包含：
            - 思考步骤内容
            - 潜在答案评估结果（用于节点选择）
            
            【参数说明】：
            parent_is_subquestion: 标识父节点是否为子问题节点，
            影响思考上下文的构建方式。
            """
            verbose_print(f"---- Generating one-step thought steps for node {self.id}...", self.verbose)

            #! ACTION: generate one-step thought step
            ost_step_list, potential_answers_list = self.generator.generate_ost_step(
                user_question=self.user_question,
                solution_trace=self.solution_trace,
                paraphrased=self.paraphrased,
                parent_is_subquestion=parent_is_subquestion,
            )
            for ost_step, potential_answers in zip(ost_step_list, potential_answers_list):
                self.children.append(
                    Reasoning_MCTS_Node(
                        parent=self,
                        depth=self.depth + 1,
                        node_type=Node_Type.OST_STEP,
                        ost_step=ost_step,
                        potential_answers=deepcopy(potential_answers),
                    )
                )

        #! create children - 根据节点类型执行相应的推理动作
        if self.node_type is Node_Type.USER_QUESTION:
            # 用户问题节点：支持A1, A2, A3, A5四种动作
            # A1: Propose an one-step thought.
            if not self.disable_a1:
                do_action_generate_ost_step()

            # A2: Propose the remaining thought steps
            do_action_generate_direct_answers()

            # A3: Propose next sub-question along with its answer.
            do_action_generate_subquestions()

            # A5: Rephrase the question/sub-question.
            if not self.disable_a5:
                do_action_generate_rephrased_user_question()

        elif self.node_type is Node_Type.REPHRASED_USER_QUESTION:
            # 重述用户问题节点：支持A1, A2, A3三种动作（不再重述）
            # A1: Propose an one-step thought.
            if not self.disable_a1:
                do_action_generate_ost_step()

            # A2: Propose the remaining thought steps
            do_action_generate_direct_answers()

            # A3: Propose next sub-question along with its answer.
            do_action_generate_subquestions()
        elif self.node_type is Node_Type.DIRECT_ANSWER:
            # 直接答案节点：终端节点，不能创建子节点
            raise ValueError("DIRECT_ANSWER node cannot create children!!")
        elif self.node_type is Node_Type.SUBQUESTION:
            # 子问题节点：支持A1, A2, A3, A4四种动作
            # A1: Propose an one-step thought.
            if not self.disable_a1:
                do_action_generate_ost_step(parent_is_subquestion=True)

            # A2: Propose the remaining thought steps
            do_action_generate_direct_answers()

            # A3: Propose next sub-question along with its answer.
            do_action_generate_subquestions()

            # A4: Answer the sub-question again.
            do_action_generate_re_subanswers()
        elif self.node_type is Node_Type.RE_SUBANSWER:
            # 重新回答节点：支持A1, A2, A3三种动作（不再重新回答）
            # A1: Propose an one-step thought.
            if not self.disable_a1:
                do_action_generate_ost_step(parent_is_subquestion=True)

            # A2: Propose the remaining thought steps
            do_action_generate_direct_answers()

            # A3: Propose next sub-question along with its answer.
            do_action_generate_subquestions()
        elif self.node_type is Node_Type.OST_STEP:
            # 一步思考节点：支持A1, A2两种动作（继续思考或给出答案）
            # A1: Propose an one-step thought.
            if not self.disable_a1:
                do_action_generate_ost_step()

            # A2: Propose the remaining thought steps
            do_action_generate_direct_answers()

        # 确保至少生成了一个子节点
        assert self.children
        return self.children

    def is_valid_leaf_node(self):
        """
        判断节点是否为有效的叶子节点
        
        【功能说明】：
        判断当前节点是否可以作为搜索树的叶子节点。
        有效的叶子节点是那些可以提供完整解决方案的终端节点。
        
        【有效叶子节点类型】：
        1. DIRECT_ANSWER: 直接答案节点，总是有效叶子节点
        2. SUBQUESTION: 子问题节点，但必须是终端子问题
           - 终端子问题：通过reach_terminal_subquestion()判断
           - 通常包含"Now we can answer"或回到原问题核心
        
        【应用场景】：
        - MCTS搜索中的终止条件判断
        - 解决方案节点的筛选
        - 搜索深度控制
        
        Returns:
            bool: 如果是有效叶子节点返回True，否则返回False
        """
        #! a valid solution can only be in SUBQUESTION type or DIRECT_ANSWER type
        return (
            self.node_type is Node_Type.SUBQUESTION and reach_terminal_subquestion(self.subquestion, self.user_question)
        ) or self.node_type is Node_Type.DIRECT_ANSWER

    def is_valid_solution_node(self):
        """
        判断节点是否为有效的解决方案节点
        
        【功能说明】：
        判断当前节点是否可以作为问题的有效解决方案。
        相比is_valid_leaf_node()，此方法包含更多的解决方案类型。
        
        【有效解决方案节点类型】：
        1. DIRECT_ANSWER: 直接答案节点，总是有效解决方案
        2. SUBQUESTION: 终端子问题节点
           - 必须通过reach_terminal_subquestion()验证
        3. OST_STEP: 终端思考步骤节点
           - 必须通过reach_terminal_ost_step()验证
           - 通常包含完整的推理结论
        
        【与is_valid_leaf_node的区别】：
        - is_valid_leaf_node: 用于MCTS搜索的终止判断
        - is_valid_solution_node: 用于解决方案的收集和评估
        
        【应用场景】：
        - 解决方案节点的识别和收集
        - 最佳解决方案的选择
        - 推理轨迹的完整性验证
        
        Returns:
            bool: 如果是有效解决方案节点返回True，否则返回False
        """
        #! a valid solution can only be in SUBQUESTION type or DIRECT_ANSWER type or OST_STEP type
        return (
            (
                self.node_type is Node_Type.SUBQUESTION
                and reach_terminal_subquestion(self.subquestion, self.user_question)
            )
            or (self.node_type is Node_Type.OST_STEP and reach_terminal_ost_step(self.ost_step))
            or self.node_type is Node_Type.DIRECT_ANSWER
        )

    def set_potential_score(self, score: float):
        """
        设置节点的潜在分数
        
        【功能说明】：
        为节点设置潜在分数，用于指导MCTS搜索的节点选择。
        潜在分数反映了节点对最终解决方案的预期贡献度。
        
        【潜在分数的作用】：
        1. 节点选择：在UCT公式中作为额外的启发信息
        2. 搜索引导：优先探索高潜在分数的节点
        3. 质量评估：评估中间节点的解决方案质量
        
        【计算来源】：
        潜在分数通常来自于对节点潜在答案的评估，
        通过模拟完成推理过程来估计节点的价值。
        
        Args:
            score (float): 节点的潜在分数值
        """
        self.potential_score = score

    def find_children(self, rollout_id: int):
        """
        查找或创建子节点（MCTS扩展阶段）
        
        【功能说明】：
        这是MCTS算法扩展阶段的核心方法，负责为当前节点生成子节点。
        如果子节点已存在则直接返回，否则调用_create_children()创建新子节点。
        
        【执行流程】：
        1. 检查是否已有子节点，如无则创建
        2. 为所有子节点设置rollout_id标识
        3. 验证子节点存在性并返回
        
        【rollout_id的作用】：
        - 追踪节点创建的rollout轮次
        - 用于调试和分析搜索过程
        - 帮助理解节点的生成时机
        
        【与_create_children的关系】：
        - find_children: 公共接口，提供缓存和ID管理
        - _create_children: 内部实现，执行具体的子节点生成逻辑
        
        Args:
            rollout_id (int): 当前rollout轮次的标识符
            
        Returns:
            List[Reasoning_MCTS_Node]: 子节点列表
        """
        self.children = self.children or self._create_children()
        for child in self.children:
            child.set_rollout_id(rollout_id)
        assert self.children
        return self.children

    def is_terminal(self):
        """
        判断节点是否为终端节点（MCTS终止条件）
        
        【功能说明】：
        判断当前节点是否应该终止MCTS搜索。
        终端节点不会再进行扩展，直接进入模拟阶段。
        
        【终端条件】：
        1. 深度限制：节点深度达到最大允许深度
           - 防止搜索过深导致的性能问题
           - 控制推理复杂度在合理范围内
        2. 有效叶子节点：节点已经是有效的解决方案
           - 直接答案节点：已获得最终答案
           - 终端子问题节点：已完成问题分解
        
        【与is_valid_leaf_node的关系】：
        - is_terminal: MCTS搜索的终止判断
        - is_valid_leaf_node: 解决方案的有效性判断
        
        【应用场景】：
        - MCTS搜索中的终止条件检查
        - 控制搜索深度和复杂度
        - 优化搜索效率
        
        Returns:
            bool: 如果是终端节点返回True，否则返回False
        """
        return self.depth >= self.max_depth_allowed or self.is_valid_leaf_node()

    def calculate_reward(self):
        """
        计算节点奖励值（MCTS模拟阶段）
        
        【功能说明】：
        计算节点在MCTS模拟阶段的奖励值，用于反向传播更新节点统计。
        奖励值反映了节点作为解决方案的质量和可信度。
        
        【奖励计算规则】：
        1. 有效叶子节点：返回节点价值（node_value）
           - 直接答案节点：答案的置信度
           - 终端子问题节点：子答案的置信度
        2. 非有效叶子节点：返回0
           - 中间节点：不直接提供解决方案
           - 无效节点：不符合解决方案要求
        
        【node_value的含义】：
        - 对于答案节点：投票选择的置信度（0-1之间）
        - 反映了答案的一致性和可靠性
        - 高置信度意味着多个候选答案趋于一致
        
        【应用场景】：
        - MCTS反向传播阶段的奖励计算
        - 节点质量评估和排序
        - 搜索策略的优化指导
        
        Returns:
            float: 节点的奖励值，有效叶子节点返回node_value，否则返回0
        """
        if self.is_valid_leaf_node():
            assert self.node_value is not None, breakpoint()
            return self.node_value
        else:
            return 0

    def skip_backprop(self):
        """
        判断是否跳过反向传播（MCTS反向传播控制）
        
        【功能说明】：
        决定当前节点是否应该跳过MCTS反向传播阶段的统计更新。
        某些节点类型不需要或不适合进行反向传播更新。
        
        【跳过反向传播的节点类型】：
        1. USER_QUESTION（用户问题节点）：
           - 作为根节点，不需要统计更新
           - 其价值不依赖于子节点的表现
        2. REPHRASED_USER_QUESTION（重述用户问题节点）：
           - 作为问题重述节点，本身不提供解决方案
           - 其价值主要体现在为后续推理提供不同视角
        
        【不跳过反向传播的节点类型】：
        - DIRECT_ANSWER: 需要统计更新以评估答案质量
        - SUBQUESTION: 需要统计更新以评估子问题分解的有效性
        - RE_SUBANSWER: 需要统计更新以评估重新回答的质量
        - OST_STEP: 需要统计更新以评估思考步骤的有效性
        
        【应用场景】：
        - MCTS反向传播阶段的控制逻辑
        - 优化搜索效率，避免无意义的统计更新
        - 确保搜索策略的合理性
        
        Returns:
            bool: 如果应该跳过反向传播返回True，否则返回False
        """
        return self.node_type is Node_Type.USER_QUESTION or self.node_type is Node_Type.REPHRASED_USER_QUESTION


def search_for_answers(args, user_question: str, question_id: int, gt_answer: str, generator: Generator):
    """
    使用MCTS搜索问题答案的主函数
    
    执行完整的MCTS搜索过程来解决数学推理问题：
    1. 构建MCTS搜索器和根节点
    2. 执行多轮rollout，每轮包含选择、扩展、模拟、反向传播
    3. 在每轮后选择当前最佳解决方案
    4. 保存搜索树和解决方案轨迹
    
    Args:
        args: 命令行参数对象，包含MCTS配置
        user_question (str): 用户问题
        question_id (int): 问题ID，用于文件命名
        gt_answer (str): 标准答案，用于评估
        generator (Generator): 推理生成器
        
    Returns:
        tuple: (模型解决方案列表, 最后rollout的ID, 所有解决方案列表)
    """
    verbose_print(
        f"********************* Searching for answers to question {question_id} ********************* ", args.verbose
    )

    #! build an MCTS searcher - 构建MCTS搜索器
    mcts_searcher = MCTS_Searcher(
        exploration_weight=args.mcts_exploration_weight,    # UCT探索权重
        weight_scheduler=args.mcts_weight_scheduler,        # 权重调度策略
        num_rollouts=args.num_rollouts,                     # 总rollout轮数
        discount=args.mcts_discount_factor,                 # 奖励折扣因子
        verbose=args.verbose,
    )

    #! build the MCTS tree - 构建MCTS树的根节点
    root_node = Reasoning_MCTS_Node(
        parent=None,
        depth=0,
        node_type=Node_Type.USER_QUESTION,
        verbose=args.verbose,
        generator=generator,
        disable_a5=args.disable_a5,
        user_question=user_question,
        expected_answer=gt_answer,
        max_depth_allowed=args.max_depth_allowed,
        disable_a1=args.disable_a1,
        enable_potential_score=args.enable_potential_score,
    )

    # 记录每轮的解决方案和统计信息
    model_solutions = []      # 每轮的最佳解决方案
    model_all_solutions = []  # 每轮的所有解决方案
    model_rollout_nodes = []  # 每轮的rollout终端节点
    
    # ================ 核心MCTS搜索循环 ================
    # 使用verbose统一控制可视化：既显示进度条又显示详细信息
    for i in (pbar := trange(args.num_rollouts, disable=not args.verbose, position=0, desc="MCTS Rollouts")):
        if args.verbose:
            print(f"\n=== MCTS Rollout {i+1}/{args.num_rollouts} ===")
            print(f"问题: {user_question[:100]}{'...' if len(user_question) > 100 else ''}")
        
        # 执行一轮MCTS rollout（选择-扩展-模拟-反向传播）
        rollout_node = mcts_searcher.do_rollout(root_node, i)
        model_rollout_nodes.append(rollout_node)

        # 从当前搜索树中选择最佳解决方案
        _, best_solution, _, chosen_node, all_solution_nodes, all_solutions = stochastic_find_best_solution(
            root_node, generator.evaluator, enable_potential_score=args.enable_potential_score
        )
        model_solutions.append(best_solution)
        model_all_solutions.append(all_solutions)
        
        # 显示当前rollout的结果和更新进度条
        if args.verbose:
            print(f"Rollout {i+1} 完成:")
            print(f"  - 找到解节点数量: {len(all_solution_nodes)}")
            print(f"  - 当前最佳解: {best_solution[:100] if best_solution else 'None'}{'...' if best_solution and len(best_solution) > 100 else ''}")
            print(f"  - 总API调用次数: {generator.io.call_counter}")
            print(f"  - 总Token使用: {generator.io.token_counter}")
            
            # 更新进度条描述以显示实时统计
            pbar.set_description(f"MCTS Rollouts (解节点: {len(all_solution_nodes)}, API调用: {generator.io.call_counter})")
        else:
            # 非verbose模式下也更新进度条描述
            pbar.set_description(f"MCTS Rollouts (解节点: {len(all_solution_nodes)}, API调用: {generator.io.call_counter})")

    # ================ 保存搜索结果和轨迹 ================
    # 保存搜索树结构（如果启用）
    if args.save_tree:
        with open(
            os.path.join(
                args.answer_sheets_dir,
                f"Question {question_id:04d} - Rollout {i+1}.tree",
            ),
            "w",
        ) as f:
            print_tree_from_root(
                mcts_searcher=mcts_searcher,
                rollout_id=i,
                root_node=root_node,
                chosen_node=chosen_node,
                file=f,
            )

    # 保存最终的所有解决方案轨迹
    js = [{"trace": node.solution_trace, "rollout_id": node.rollout_id} for node in all_solution_nodes]
    with open(os.path.join(args.answer_sheets_dir, f"Question {question_id:04d} - Final Solutions.json"), "w") as f:
        json.dump(js, f)

    # 保存每轮rollout的解决方案轨迹
    js2 = [{"trace": node.solution_trace, "rollout_id": i} for i, node in enumerate(model_rollout_nodes)]
    with open(os.path.join(args.answer_sheets_dir, f"Question {question_id:04d} - Rollout Solutions.json"), "w") as f:
        json.dump(js2, f)

    # 保存潜在答案历史（如果启用潜在分数评估）
    if args.enable_potential_score:
        js = [node.potential_answers_history for node in all_solution_nodes]
        with open(os.path.join(args.answer_sheets_dir, f"Question {question_id:04d} - Potentials.json"), "w") as f:
            json.dump(js, f)

    return model_solutions, i, model_all_solutions
